import os
import os.path
import kuai

pin, pout, perr = os.popen3(os.path.join(kuai.DFF_BIN, "dffjob.exe"))

def get_dff_types(mol):
    molfile = kuai.format_mol(mol)
    s = "#DFF: Typing DFF\n"
    s += 'mol = dffpy.parse_mol("""\n'
    s += molfile
    s += '""", "msd")\n'
    s += "print 'Job Finished at', datetime.now()\n"
    pin.write(s)
    pin.flush()

    types = []
    while True:
        i = pout.readline()
        if i.strip().startswith('Job Finished at'):
            break
        else:
            types.append(i.strip())
    return types

def assign_dff_types(mol):
    types = get_dff_types(mol)
    assert len(types) == len(mol.atoms)
    for i in range(len(types)):
        mol.atoms[i].type = types[i]
    return mol

def get_struct_types(mol):
    molfile = kuai.format_mol(mol)
    s = "#DFF: Typing Structure\n"
    s += 'mol = dffpy.parse_mol("""\n'
    s += molfile
    s += 'Shit, "msd")\n'
    s += "print 'Job Finished at', datetime.now()\n"
    # print type(s), type(pin), pin
    print s
    pin.write(s)
    pin.flush()

    types = []
    while True:
        i = pout.readline()
        print i
        if i.strip().startswith('Job Finished at'):
            break
        else:
            types.append(i.strip())
    print types
    return types

def assign_struct_types(mol):
    types = get_struct_types(mol)
    assert len(types) == len(mol.atoms)
    for i in range(len(types)):
        mol.atoms[i].type = types[i]
    return mol


mol = kuai.read_mol(r"F:/Madlee/Job/TEAM/MayBridge/drug.msd")
print mol.formula()
ff = kuai.format_mol(mol)
print ff
types = get_struct_types(mol)
print types
# write_mol("xxx.msd", mol)
